A comment here about the shape of count is useful (since it's very specific) and commonplace in the code: https://github.com/google-deepmind/optax/blob/main/optax/_src/transform.py#L88

Oh nice. I see it is

count: chex.Array  # shape=(), dtype=jnp.int32.

everywhere.

A brief description of AdEMAMix and how it actually operates may be useful. For example, a comparison to Adam (on which it is based).

Alternatively, pseudo-code resembling https://github.com/google-deepmind/optax/blob/main/optax/_src/alias.py#L586-L594 may be helpful.

This is essentially duplicate information relative to "Exponential decay rate". I'd suggest mentioning EMA in a general description above, and then making the arg information resemble that of AdamW.

Why not make alpha (and b3) a base.ScalarOrSchedule instead?

I noticed the patter with the majority of the other classes that did not use base.ScalarOrSchedule as the type and tried to replicate that.

Of course, the authors in the paper specifically recommend scheduling b3 and alpha because of early training instabilities, so this is a bit of a different use case.

I will make this change.

I'm not sure why you have multiple counts. Would just the first suffice?

So the first count replicates the same entitiy in adam. The second is for keeping track of the second momentum, so the user might want to do different things with each.

Does that make sense?

I will make an inline comment to this fact

I'm still not convinced that there's utility in having separate counts - they're incremented identically in the actual code. I would lean towards having a single count, but I don't think this is blocking.

Super cool work - reproducing results is definitely under-appreciated. This is really nice!

I know this isn't done elsewhere in optax, but generally I prefer opaque constants (e.g. 1 in this case) to be passed in via kwargs. Consider this optional though.

so:

m2 = otu.tree_update_moment(updates, state.m2, c_b3, order=1)

is what you are describing?

Exactly.

I think the convention in this library is that imports from python generally go at the top, so

from typing import...

import chex
...

Ah! Good point. (I was going to make a meta point about this repo not having any pre-commit hooks but...).

Could you update this docstring too? The ademamix docstring looks great. You may simply refer here to the docstring of ademamix for e.g. the pseudocode and copy-paste e.g. the descriptions of the arguments.

So to clarify: just a link to the ademamix docstring? That is, don't follow the convention of:

  Description

  References:

  Args:

  Returns:
    A `GradientTransformation` object.

?

Note: the above convention is prevalent throughout optax/_src/transform.py for similar scale_by_X functions.

You can use the following docstring

  """Scale updates according to the Ademamix algorithm.

  See :func:`optax.contrib.ademamix.` for a full description of the algorithm.

  References:
    Pagliardini et al, `The AdEMAMix Optimizer: Better, Faster, Older
    <https://arxiv.org/abs/2409.03137>`_, 2024

  Args:
    learning_rate: A global scaling factor, either fixed or evolving along
      iterations with a scheduler, see :func:`optax.scale_by_learning_rate`.
    b1: Exponential decay rate to track the fast EMA.
    b2: Exponential decay rate to track the second moment of past gradients.
    b3: Exponential decay rate to track the slow EMA.
    alpha: Mixing coefficient in the linear combination fo the fast and 
      slow EMAs. 
    eps: A small constant applied to denominator outside of the square root
      (as in the Adam paper) to avoid dividing by zero when rescaling.
    eps_root: A small constant applied to denominator inside the square root (as
      in RMSProp), to avoid dividing by zero when rescaling. This is needed for
      instance when computing (meta-)gradients through Adam.
    mu_dtype: Optional `dtype` to be used for the first order accumulator; if
      `None` then the `dtype` is inferred from `params` and `updates`.

  Returns:
    The corresponding `GradientTransformation`.
  """

Could you add a description of the attributes in the docstring?
See e.g.

I took a stab. LMK if there are better descriptions I could do. I (theoretically) could add the math update if you believe that would be warranted. I also want to get the description right. nu is still the estimate of the second moment, but now the first moment is a combination of m1 and m2 so I just stated their respective EMA types.

Looks great to me.
I think I would say "fast EMA of first moment" rather than first moment personally (same for "slow EMA" and this would apply in the other docstrings). But if you have stronger preferences I won't argue against.

You can use the following. (the current formatting won't be well handled by docs).

.. note::

    AdEMAMix consists in leveraging very old gradients. Therefore,
    the method is best suited to settings where the number of iterations is
    important. The paper reports on this effect in Appendix C.1.5, showing how
    smaller values of b3 (e.g. b3 = 0.999) can be better for low iterations
    scenarios. Moreover, retaining gradient information over many thousands of
    steps can pose a problem in domains requiring fast adaptation to a sudden
    distribution shift, or general cases in which the distribution is
    non-stationary.

Also, you may build the docs using (from the folder docs)

make html-noplot

that way you can check whether the docs are well formatted.

You can use the following docstring

  """Scale updates according to the Ademamix algorithm.

  See :func:`optax.contrib.ademamix.` for a full description of the algorithm.

  References:
    Pagliardini et al, `The AdEMAMix Optimizer: Better, Faster, Older
    <https://arxiv.org/abs/2409.03137>`_, 2024

  Args:
    learning_rate: A global scaling factor, either fixed or evolving along
      iterations with a scheduler, see :func:`optax.scale_by_learning_rate`.
    b1: Exponential decay rate to track the fast EMA.
    b2: Exponential decay rate to track the second moment of past gradients.
    b3: Exponential decay rate to track the slow EMA.
    alpha: Mixing coefficient in the linear combination fo the fast and 
      slow EMAs. 
    eps: A small constant applied to denominator outside of the square root
      (as in the Adam paper) to avoid dividing by zero when rescaling.
    eps_root: A small constant applied to denominator inside the square root (as
      in RMSProp), to avoid dividing by zero when rescaling. This is needed for
      instance when computing (meta-)gradients through Adam.
    mu_dtype: Optional `dtype` to be used for the first order accumulator; if
      `None` then the `dtype` is inferred from `params` and `updates`.

  Returns:
    The corresponding `GradientTransformation`.
  """

Typo: Exponential

Looks great to me.
I think I would say "fast EMA of first moment" rather than first moment personally (same for "slow EMA" and this would apply in the other docstrings). But if you have stronger preferences I won't argue against.